from keras.applications.vgg16 import VGG16
from keras.models import Sequential
from keras.layers import Dropout, Flatten, Dense
from config import Config

vgg16_model = VGG16(weights='imagenet', include_top=False, input_shape=Config.CLASS_TARGET_SIZE)


# 搭建新的全连接层
top_model = Sequential()
top_model.add(Flatten(input_shape=vgg16_model.output_shape[1:]))
top_model.add(Dense(256, activation='relu'))
top_model.add(Dropout(0.5))
top_model.add(Dense(Config.CLASS_NUMBER, activation='softmax'))

model = Sequential()
model.add(vgg16_model)
model.add(top_model)